package limiter

import (
	"context"
	"fmt"
	"github.com/redis/go-redis/v9"
	"time"
)

// RedisLimiter 基于Redis的分布式限流器
type RedisLimiter struct {
	client *redis.Client
}

func NewRedisLimiter(client *redis.Client) *RedisLimiter {
	return &RedisLimiter{client: client}
}

func (r *RedisLimiter) Allow(ctx context.Context, key string, limit int, window time.Duration) (bool, error) {
	now := time.Now().UnixNano()
	windowMicro := window.Nanoseconds()

	luaScript := `
	local key = KEYS[1]
	local now = tonumber(ARGV[1])
	local window = tonumber(ARGV[2])
	local limit = tonumber(ARGV[3])
	
	redis.call('ZREMRANGEBYSCORE', key, 0, now - window)
	
	local current = redis.call('ZCARD', key)
	
	if current >= limit then
		return 0
	end
	
	redis.call('ZADD', key, now, now)
	redis.call('EXPIRE', key, window/1000000000)
	return 1
	`

	res, err := r.client.Eval(ctx, luaScript, []string{key}, now, windowMicro, limit).Result()
	if err != nil {
		return false, fmt.Errorf("redis eval failed: %v", err)
	}

	allowed, ok := res.(int64)
	if !ok {
		return false, fmt.Errorf("invalid redis response")
	}

	return allowed == 1, nil
}
